#!/usr/bin/env python2
#-*- coding: utf-8 -*-

import os
import sys
import socket
import time
import subprocess
import fcntl
import types
import errno
import getopt

from config import Config
from storage import Storage
from utils import Exp, _dmsg, _dwarn, _derror, _get_value, _exec_shell
from metanode_balance_with_diskmap import get_lst_v4
from metadata_balance_with_diskmap import get_new_locs
from diskmap import name2map, map2name, name2site

config = Config()
storage = Storage(config)

ZONE_TYPES = ["site", "zone", "node"]
LICH_REPLICA_MAX  = 4

def tmpfortest(chunks):
    for chunk in chunks:
        new_locs = ["wu2/0", "wu3/0"]
        print 'new locs', new_locs
        chunk["disks"] = new_locs


def fetch_job(job):
    #chunk = {"chkid": xxx, "disks": ["wu2/0", "wu3/0"]}
    #jobs = [chunk, chunk, ...]
    print 'job', job
    chkid = job["chkid"]
    disks = ",".join([x.split("/")[0] for x in job["disks"]])
    cmd = "%s --chunkmove %s %s" % (config.inspect, chkid, disks)

    retry = 0
    retry_max = 10
    while True:
        try:
            _exec_shell(cmd)
            break
        except Exp, e:
            if retry > retry_max:
                raise

            if e.errno not in [errno.EBUSY, errno.EAGAIN, errno.ETIMEDOUT, errno.ENONET]:
                raise

            _dwarn(e)
            time.sleep(3)
            retry = retry + 1

def _fetch_jobs(jobs, chunks):
    succ = 0
    fail = 0
    total = len(jobs)
    errors = {}

    for job in jobs:
        for chunk in chunks:
            if chunk["chkid"] == job["chkid"]:
                break

        if sorted(job["disks"]) == sorted(chunk["disks"]):
            succ = succ + 1
            print '%s no need to move' % (job)
            continue

        try:
            fetch_job(job)
            succ = succ + 1
        except Exp, e:
            errors.update({job["chkid"]: str(e)})
            fail = fail + 1
            pass

    assert(succ + fail == total)

    if fail:
        raise Exp(errno.EAGAIN, str(errors))

def fetch_jobs(jobs, chunks):
    retry = 0
    retry_max = 3
    while True:
        try:
            _fetch_jobs(jobs, chunks)
            break
        except Exp, e:
            if retry > retry_max:
                raise

            if e.errno not in [errno.EBUSY, errno.EAGAIN, errno.ETIMEDOUT, errno.ENONET]:
                raise

            _dwarn(e)
            time.sleep(3)
            retry = retry + 1

def check(jobs, ablenodes):
    for i in range(LICH_REPLICA_MAX):
        print 'replica: ', i
        disks = []
        for x in jobs:
            if (i+1) > len(x["disks"]):
                break
            disks.append(x["disks"][i])
        keys = sorted(list(set(disks)))
        for k in keys:
            print 'k', k, disks.count(k)

def get_jobs(ablenodes, chunks):
    jobs = []

    for chunk in chunks:
        chunk_ = {"chkid": chunk["chkid"]}
        disks_rep = []
        #disks_rep = [{"replica": "", "disk": "disk"}, {""}]
        for a in ablenodes:
            for chunk_rep in a["chunk_reps"]:
                if chunk_rep["chkid"] == chunk["chkid"]:
                    disk_rep = {"replica": chunk_rep["replica"], "name": a["name"]}
                    disks_rep.append(disk_rep)

        disks = []
        for i in range(LICH_REPLICA_MAX):
            find = False
            for d in disks_rep:
                if d["replica"] == i:
                    assert(len(disks) == i)
                    disks.append(d["name"])
                    find = True
                    break

            if not find:
                break

        assert(len(disks) >= 2)
        
        chunk_.update({"disks": disks})
        jobs.append(chunk_)
        
    return jobs

def get_ablenodes(lst):
    #lst = {'wu4/0': 'admin', 'wu2/0': 'meta', 'wu3/0': 'meta'}
    #ablenode = {"site": "site", "zone": "xxx", "node": xxx, "chunk_reps": [chunk_rep]}
    #ablenodes = [ablenode, ablenode, ...]
    ablenodes = []
    for d in lst.keys():
        dmap = name2map(d) 
        name = map2name(dmap)
        ablenode = {"name": name, "site": dmap[0], "zone": dmap[1], "node": dmap[2], "chunk_reps": []}
        ablenodes.append(ablenode)
    return ablenodes

def multpath_chunks_get():
    rs = storage.find("/iscsi/", "vol")
    chkinfos = [storage.chkinfo(x) for x in rs]
    chkinfos_multpath = [x for x in chkinfos if storage.ismultpath(x["chkid"])]
    #print chkinfos_multpath

    for x in chkinfos_multpath:
        disks = x["disks"]
        _disks = []
        for d in disks:
            _disk = "%s/0" % (d.split(":")[0])
            _disks.append(_disk)
        x["disks"] = _disks

    #files = [x for x in rs if storage.isfile(x)]
    #multpath_files = [x for x in files if storage.ismultpath(x)]
    #chkids = [storage.chkid(x) for x in multpath_files]
    #print 'rs', rs
    #print 'files', files
    #print 'multpath', multpath_files
    #print chkinfos_multpath
    return chkinfos_multpath

def all_chunks_get():
    rs = storage.find("/", skips=["/system/unlink"])

    chkinfos = []
    for x in rs:
        try:
            chkinfos.append(storage.chkinfo(x))
        except Exp, e:
            _dwarn('get chunkinfo fail: %s, %s' % (x, e))

    chkinfos = [storage.chkinfo(x) for x in rs]

    for x in chkinfos:
        disks = x["disks"]
        _disks = []
        for d in disks:
            _disk = "%s/0" % (d.split(":")[0])
            _disks.append(_disk)
        x["disks"] = _disks

    return chkinfos

def has_chkid(ablenodes, chkid):
    for a in ablenodes:
        chkids = [x["chkid"] for x in a["chunk_reps"]]
        if chkid in chkids:
            return True

    return False

def get_sites_from_lst(lst):
    #lst = {'wu4/0': 'admin', 'wu2/0': 'meta', 'wu3/0': 'meta'}
    #sites = ["site1", "site2"]
    sites = []
    for d in lst.keys():
        site = name2site(d)
        if site not in sites:
            sites.append(site)
    return sites

def ablenodes_init(ablenodes, chunks, replica):
    for chunk in chunks:
        if replica + 1 > len(chunk["disks"]):
            continue

        for n in ablenodes:
            name = chunk["disks"][replica]
            name = map2name(name2map(name))
            chunk_rep = {"chkid": chunk["chkid"], "replica": replica, "name": name}
            n2 = map2name(name2map(n["name"]))
            if name == n2:
                n["chunk_reps"].append(chunk_rep)
                break

def get_sum_weight_ablenodes(ablenodes, replica):
    #print 'get_sum_wight_ablendoes', ablenodes
    sum_weight = 0
    for a in ablenodes:
        #print 'a', a
        reps = [rep for rep in a["chunk_reps"]]
        reps = [x for x in reps if x["replica"] == replica]
        num = len(reps)
        sum_weight = sum_weight + num

    return sum_weight

def get_max_ablezone(ablenodes, replica):
    zones = get_zones(ablenodes)
    #print 'zones', zones

    zones_able = []
    for zone in zones:
        ables = []
        for a in ablenodes:
            _zone = "%s.%s" % (a["site"], a["zone"])
            if _zone == zone:
                ables.append(a)
        zones_able.append(ables)

    max_ = zones_able[0]
    for x in zones_able:
        num_x = get_sum_weight_ablenodes(x, replica)
        num_max = get_sum_weight_ablenodes(max_, replica)

        if num_x > num_max:
            max_ = x

    return max_

def get_min_ablezone(ablenodes, replica):
    zones = get_zones(ablenodes)
    zones_able = []
    for zone in zones:
        ables = []
        for a in ablenodes:
            _zone = "%s.%s" % (a["site"], a["zone"])
            if _zone == zone:
                ables.append(a)
        zones_able.append(ables)

    min_ = zones_able[0]
    for x in zones_able:
        num_x = get_sum_weight_ablenodes(x, replica)
        num_min = get_sum_weight_ablenodes(min_, replica)

        if num_x < num_min:
            min_ = x

    return min_

def get_max_ablesite(ablenodes, replica):
    sites = get_sites(ablenodes)
    sites_able = []
    for site in sites:
        ables = []
        for a in ablenodes:
            if a["site"] == site:
                ables.append(a)
        sites_able.append(ables)

    max_ = sites_able[0]
    for x in sites_able:
        num_x = get_sum_weight_ablenodes(x, replica)
        num_max = get_sum_weight_ablenodes(max_, replica)

        if num_x > num_max:
            max_ = x

    return max_

def get_min_ablesite(ablenodes, replica):
    sites = get_sites(ablenodes)
    sites_able = []
    for site in sites:
        ables = []
        for a in ablenodes:
            if a["site"] == site:
                ables.append(a)
        sites_able.append(ables)

    min_ = sites_able[0]
    for x in sites_able:
        num_x = get_sum_weight_ablenodes(x, replica) 
        num_min = get_sum_weight_ablenodes(min_, replica)

        if num_x < num_min:
            min_ = x

    return min_

def get_max_ablenode(ablenodes, replica):
    max_ = ablenodes[0]

    for x in ablenodes:
        num_x = get_sum_weight_ablenodes([x], replica)
        num_max = get_sum_weight_ablenodes([max_], replica)

        if num_x > num_max:
            max_ = x

    return max_

def get_min_ablenode(ablenodes, replica):
    min_ = ablenodes[0]

    for x in ablenodes:
        num_x = get_sum_weight_ablenodes([x], replica) 
        num_min = get_sum_weight_ablenodes([min_], replica)

        if (num_x < num_min):
            min_ = x

    return min_

def get_sites(ablenodes):
    sites = []
    for a in ablenodes:
        site = a["site"]
        if site not in sites:
            sites.append(site)

    return sites

def get_zones(ablenodes, parent=None):
    zones = []

    #print 'ablenodes', ablenodes
    for a in ablenodes:
        site = a["site"]
        zone = "%s.%s" % (a["site"], a["zone"])
        if parent is not None:
            if (parent == site) and (zone not in zones):
                zones.append(zone)
        else:
            if zone not in zones:
                zones.append(zone)

    return zones

def get_ablenodes_byparent(_ablenodes, parent=None):
    ablenodes = []

    for a in _ablenodes:
        site = a["site"]
        zone = a["zone"]
        name = a["name"]
        if parent is not None:
            if (len(parent.split(".")) == 1):
                _parent = "%s" % (site)
            else:
                assert(len(parent.split(".")) == 2)
                _parent = "%s.%s" % (site, zone)

            if (parent == _parent) and (name not in [x["name"] for x in ablenodes]):
                ablenodes.append(a)
        else:
            if (name not in [x["name"] for x in nodes]):
                ablenodes.append(a)

    return ablenodes

def is_balance(max_, min_, replica, diff=2):
    #max_ = [ablenode, ablenode]
    #min_ = [ablenode, ablenode]
    #todo 只检查指定replica 是否平衡
    assert(type(max_) == type([]))
    assert(type(min_) == type([]))
    max_weight = get_sum_weight_ablenodes(max_, replica)
    min_weight = get_sum_weight_ablenodes(min_, replica)
    print 'diff', (max_weight - min_weight)
    return (max_weight - min_weight) <= diff

def ablenodes_move_max2min(max_, min_, replica):
    min_chkids = [x["chkid"] for x in min_["chunk_reps"]]
    index = -1
    chunk_rep = {}

    for i in range(len(max_["chunk_reps"])):
        _rep = max_["chunk_reps"][i]
        if (_rep["replica"] == replica) and (_rep["chkid"] not in min_chkids):
            index = i
            chunk_rep = _rep
            break
            
    if (index < 0):
        raise Exp(errno.ENOSPC, "%s has in %s" % (_rep, min_))

    assert(chunk_rep["chkid"])

    del max_["chunk_reps"][index]
    min_["chunk_reps"].append(chunk_rep)
    print 'move %s from %s to %s' % (chunk_rep, max_["name"], min_["name"]) 

def ablenodes_balance_sites(ablenodes, replica):
    warn_num = 0
    warn_max = 1000

    i = 0
    while True:
        max_ = get_max_ablesite(ablenodes, replica)
        min_ = get_min_ablesite(ablenodes, replica)

        if is_balance(max_, min_, replica):
            break

        max_node = None
        min_node = None

        for m in max_:
            for chunk_rep in m["chunk_reps"]:
                if chunk_rep["replica"] != replica:
                    continue

                if has_chkid(min_, chunk_rep["chkid"]):
                    warn_num = warn_num + 1
                    continue
                else:
                    max_node = m
                    min_node = min_[0]
                    break

        if (max_node is None) or (min_node is None) or (max_node["name"] == min_node["name"]):
            print 'balance sites may be ok.'
            break

        if warn_num > warn_max:
            print 'balance sites warn_num much more, may be ok. need break.'
            break

        ablenodes_move_max2min(max_node, min_node, replica)
        print 'balance site, %s' % (i)
        i = i + 1

    sites = get_sites(ablenodes)
    for site in sites:
        print 'balance nodes of site: %s' % (site)
        nodes = get_ablenodes_byparent(ablenodes, parent=site)
        ablenodes_balance_nodes(nodes, replica)

def ablenodes_balance_zones(ablenodes, replica):
    warn_num = 0
    warn_max = 1000

    i = 0
    while True:
        max_ = get_max_ablezone(ablenodes, replica)
        min_ = get_min_ablezone(ablenodes, replica)

        if is_balance(max_, min_, replica):
            break

        max_node = None
        min_node = None

        for m in max_:
            for chunk_rep in m["chunk_reps"]:
                if chunk_rep["replica"] != replica:
                    continue

                if has_chkid(min_, chunk_rep["chkid"]):
                    warn_num = warn_num + 1
                    continue
                else:
                    max_node = m
                    min_node = min_[0]
                    break


        if (max_node is None) or (min_node is None) or (max_node["name"] == min_node["name"]):
            print 'balance zones may be ok.'
            break

        if warn_num > warn_max:
            print 'balance zons warn_num much more, may be ok. need break.'
            break

        ablenodes_move_max2min(max_node, min_node, replica)
        print 'balance zone, %s' % (i)
        i = i + 1

    zones = get_zones(ablenodes)
    #print zones
    #print ablenodes
    for zone in zones:
        print 'balance nodes of zone: %s' % (zone)
        nodes = get_ablenodes_byparent(ablenodes, parent=zone)
        ablenodes_balance_nodes(nodes, replica)

def ablenodes_balance_nodes(ablenodes, replica):
    warn_num = 0
    warn_max = 1000

    i = 0
    while True:
        max_ = get_max_ablenode(ablenodes, replica)
        min_ = get_min_ablenode(ablenodes, replica)
        #print 'max_', max_
        #print 'min_', min_
        #print 'replica_', replica
        #print 'ablenodes', ablenodes

        if is_balance([max_], [min_], replica):
            break

        if (max_["name"] == min_["name"]):
            print 'balance nodes may be ok.'
            break

        ablenodes_min = [x for x in ablenodes if x["name"] != max_["name"]]
        ablenodes_min_sort = sorted(ablenodes_min, key=lambda x:get_sum_weight_ablenodes([x], replica), reverse=False)
        #print 'ablenodes_min_sort', ablenodes_min_sort
        print 'max', max_["name"], get_sum_weight_ablenodes([max_], replica)
        for x in ablenodes_min_sort:
            print 'min', x["name"], get_sum_weight_ablenodes([x], replica)

        found = False
        for x in ablenodes_min_sort:
            try:
                ablenodes_move_max2min(max_, x, replica)
                found = True
                break
            except Exp, e:
                if e.errno == errno.ENOSPC:
                    warn_num = warn_num + 1
                    _dwarn(e)
                    continue

        if warn_num > warn_max:
            print 'balance nodes warn_num much more, may be ok. need break.'
            break

        if not found:
            print 'balance nodes not found, may be ok.'
            break

        print 'balance node, %s' % (i)
        i = i + 1
        
def ablenodes_balance(ablenodes, replica):
    sites = get_sites(ablenodes)
    zones = get_zones(ablenodes)

    #print 'sites', sites
    #print 'zones', zones
    if len(sites) >= 3:
        ablenodes_balance_sites(ablenodes, replica)
    elif len(zones) >= 3:
        ablenodes_balance_zones(ablenodes, replica)
    else:
        ablenodes_balance_nodes(ablenodes, replica)

def metabalance(lst, chunks):
    for chunk in chunks:
        locs = chunk["disks"]
        master = chunk["disks"][0]
        print 'chunk', chunk
        new_locs = get_new_locs(lst, locs, master)
        print 'new locs', new_locs

        if sorted(locs) == sorted(new_locs):
            _dmsg("chunk %s need not move" % (chunk))
            continue

        job = {"chkid": chunk["chkid"], "disks": new_locs}
        try:
            fetch_job(job)
            chunk["disks"] = new_locs
        except Exp, e:
            _dwarn("metabalance warn: %s" % (e))
            pass

def balance(lst, chunks):
    """
    name = "wu2/0"
    disks = [name, name, ...]
    lst = {'wu4/0': 'admin', 'wu2/0': 'meta', 'wu3/0': 'meta'}
    chunk = {"chkid": xxx, "disks": ["wu2/0", "wu3/0"]}
    chunk_rep = {"chkid": xxx, "replica": x, "name": x}
    ablenode = {"name":"name", "site": "site", "zone": "xxx", "node": xxx, "chunk_reps": [chunk_rep]}
    jobs = [chunk, chunk, ...]
    chunks = [chunk, chunk, ...]
    ablenodes = [ablenode, ablenode, ...]
    """
    ablenodes = []
    jobs = []

    #make sure in differen zone
    metabalance(lst, chunks)

    #make sure avg in zone, if it is multpath
    chunks = [x for x in chunks if storage.ismultpath(x["chkid"])]
    ablenodes = get_ablenodes(lst)

    for replica in range(LICH_REPLICA_MAX):
        print 'ablenodes init, replica: %s' % (replica)
        ablenodes_init(ablenodes, chunks, replica)

    for replica in range(LICH_REPLICA_MAX):
        print 'balance replica: %s' % (replica)
        ablenodes_balance(ablenodes, replica)

    jobs = get_jobs(ablenodes, chunks)

    fetch_jobs(jobs, chunks)

    print 'old:'
    check(chunks, ablenodes)

    print 'new:'
    check(jobs, ablenodes)

    print 'iscsi_balance ok'

def balance_site(site, _lst, _chunks):
    #lst = {'wu4/0': 'admin', 'wu2/0': 'meta', 'wu3/0': 'meta'}
    #sites = ["site1", "site2"]
    #chunk = {"chkid": xxx, "disks": ["wu2/0", "wu3/0"]}
    lst = {}
    chunks = []

    for k,v in _lst.items():
        if (site == name2site(k)):
            lst.update({k: v})

    for chunk in _chunks:
        if (site == name2site(chunk["disks"][0])):
            chunks.append(chunk)

    print 'balance site: %s, lst: %s, chunks: %s' % (site, lst, chunks)
    balance(lst, chunks)

def balance_check_site(chunks, lst):
    '''
    chunk = {"chkid": xxx, "disks": ["wu2/0", "wu3/0"], "path": xx}
    chunks = [chunk, chunk, ...]
    '''
    for chunk in chunks:
        try:
            site = storage.siteof(chunk["path"])
        except Exp, e:
            _dwarn("get siteof fail %s %s" % (chunk, e))

        if storage.site_isnull(site):
            site = name2site(chunk["disks"][0])

        disks = []
        new = 0
        for disk in chunk["disks"]:
            if (name2site(disk) == site):
                disks.append(disk)
            else:
                for d in lst.keys():
                    if (name2site(d) == site) and (d not in disks):
                        disks.append(d)
                        new = new + 1
                        break

        if (new):
            chunk['disks'] = disks
            try:
                fetch_job(chunk)
            except Exp, e:
                _dwarn("metabalance warn: %s %s" % (chunk, e))

    return chunks

def balance_all():
    '''
    chunk = {"chkid": xxx, "disks": ["wu2/0", "wu3/0"], "path": xx}
    chunks = [chunk, chunk, ...]
    '''
    lst = get_lst_v4()
    chunks = all_chunks_get()
    print 'lst: ', lst

    #make sure in the assign site
    chunks = balance_check_site(chunks, lst)

    sites = get_sites_from_lst(lst)
    for site in sites:
        balance_site(site, lst, chunks)

def main():
    print 'hi boy'

if __name__ == '__main__':
    main()
